import os
import sys

dir = sys.argv[1]
dir_path = os.path.abspath(dir)
files = os.listdir(dir_path)
py_files = [fname for fname in files if ("SelfTrain_" in fname) and (".py" in fname)]
bash_files = [os.path.join(dir_path,fname.replace(".py", ".sh")) for fname in py_files]
for idx, bash_file in enumerate(bash_files):
    with open(bash_file, 'w') as fw:
        print(bash_file)
        fw.write("\n".join([
           "#!/bin/bash",
           "#SBATCH -N 1",
           "#SBATCH --gres=gpu:1",
           "#SBATCH --mem=20G",
           "#SBATCH -p cuda10-ceshi",
           "echo $(hostname) $CUDA_VISIBLE_DEVICES\n"
        ]))
        cmd_line = "srun -N 1 -n 1 --mem=20G --output=%s --gres=gpu:1 python %s $CUDA_VISIBLE_DEVICES"%(
                            py_files[idx].replace('.py', '.log'), py_files[idx])
        fw.write(cmd_line)

